import numpy as np
import logging

class LeafNode(object):
    def __init__(self, yhat, id_leaf):
        self.yhat = yhat
        self.id_leaf = id_leaf

    def data_down_tree(self, X):
        return np.outer(np.ones(X.shape[0]), self.yhat), np.outer(np.ones(X.shape[0]), self.id_leaf)


class InternalNode(object):
    def __init__(self, dim, threshold, left_child, right_child):
        self.dim = dim
        self.threshold = threshold
        self.left_child = left_child
        self.right_child = right_child

        assert left_child is not None
        assert right_child is not None

    def data_down_tree(self, X):
        # branch masks
        if X[self.dim].dtype.type is np.float_ or X[self.dim].dtype.type is np.int_ :
            left_mask = X[self.dim] < self.threshold
        elif X[self.dim].dtype.type is np.string_  or X[self.dim].dtype.type is np.bool_:
            left_mask = X[self.dim] == self.threshold
        else:
            logging.error('dimension type not recognize: ', self.dim.dtype)
            raise Exception('ERROR: one of the columns of the data is not string nor float')
        right_mask = np.logical_not(left_mask)

        # sanity check: everybody needs to go somewhere and nobody can
        # go both directions
        assert np.all(np.logical_xor(left_mask, right_mask))

        go_left = np.any(left_mask)
        go_right = np.any(right_mask)

        # evaluate branches
        if go_left:
            Y_left, prediction_similarity_left = self.left_child.data_down_tree(X[left_mask])

        if go_right:
            Y_right, prediction_similarity_right = self.right_child.data_down_tree(X[right_mask])

        if go_left and go_right:
            # sanity check: make sure we got compatible results from the two subtrees
            assert Y_left.dtype == Y_right.dtype
            assert Y_left.shape[1] == Y_right.shape[1]

            # rebuild Y in the same order as X
            Y = np.zeros((X.shape[0], Y_left.shape[1]), dtype=Y_left.dtype)
            #pred_leaves  = np.zeros((X.shape[0], Y_left.shape[1]), dtype=Y_left.dtype)
            pred_leaves  = np.zeros((X.shape[0], 1), dtype=Y_left.dtype)

            Y[left_mask,:] = Y_left
            Y[right_mask,:] = Y_right

            pred_leaves[left_mask,:] = prediction_similarity_left
            pred_leaves[right_mask,:] = prediction_similarity_right

            return Y, pred_leaves

        elif go_left:
            return Y_left, prediction_similarity_left

        elif go_right:
            return Y_right, prediction_similarity_right

        else:
            raise Exception("This can't happen.")

